/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#ifndef __SQUARED_LOSS_HH__
#define __SQUARED_LOSS_HH__

#include "dormandprince.hh"
#include "dyadicp.hh"
#include "loss.hh"

#include <boost/optional.hpp>
#include <cmath>
#include <stdint.h>

#include "frsqrt.hh"

namespace flassol
{
  class SquaredLoss : public Loss
    {
      private:
        static LossFactory factory;

        double   sumy;
        double   sumysq;
        uint64_t n;
        uint64_t skipped_update;
        uint64_t n_smalls;
        uint64_t total_updates;

        class DFunc 
          {
            private:
              float adotb;
              float anormplusbnorm;
              float wdotx;
              float xnorm;
              float lambda;
              float y;

            public:
              DFunc (float _adotb,
                     float _anormplusbnorm,
                     float _wdotx,
                     float _xnorm,
                     float _lambda,
                     float _y)
                : adotb (_adotb),
                  anormplusbnorm (_anormplusbnorm),
                  wdotx (_wdotx),
                  xnorm (_xnorm),
                  lambda (_lambda),
                  y (_y)
                {
                }

              float
              operator() (float h,
                          float s) const
                {
                  float kappa = (fabsf (s) > 0) ? -lambda * h / s : 0.0f;

                  // TODO: rewrite dyadic::p in terms of decay not kappa

                  std::pair<bool, float> p = dyadic::p (s,
                                                        adotb,
                                                        anormplusbnorm,
                                                        wdotx,
                                                        xnorm,
                                                        kappa);

                  return p.second - y;
                }
          };

      public:
        explicit
        SquaredLoss () : sumy (0),
                         sumysq (0),
                         n (0),
                         skipped_update (0),
                         n_smalls (0),
                         total_updates (0)

          { 
          }

        virtual float 
        loss (float p,
              float y)
          {
            sumy += y;
            sumysq += y * y;
            ++n;

            float z = (p - y);
            return 0.5 * z * z;
          }

        virtual float 
        dldp (float p,
              float y)
          {
            return (p - y);
          }

        virtual std::pair<float, float>
        invariant_update (float p,
                          float wdotx,
                          float adotb,
                          float y,
                          float eta,
                          float anormplusbnorm,
                          float xnorm,
                          bool  adaptive,
                          float adaptivexnorm,
                          float lambda)
          {
            ++total_updates;

            if (p == y)
              {
                return std::make_pair (0.0f, 1.0f);
              }
            else if (anormplusbnorm == 0.0f)
              {
                float etaxtx = (adaptive) ? eta * adaptivexnorm * rsqrtf (xnorm)
                                          : eta;
                float xtx = (adaptive) ? adaptivexnorm : xnorm;

                if (fabs (etaxtx) < 1e-5)
                  {
                    return std::make_pair ((p - y) * etaxtx / xtx, 1.0f);
                  }
                else
                  {
                    return 
                      std::make_pair ((p - y) * (1.0f - expf (-etaxtx)) / xtx,
                                      1.0f);
                  }
              }
            else
              {
                // TODO: for quantile and hinge loss, the dynamics
                // stop when you hit the label.  this includes regularization,
                // so the user can specify a massive importance weight and
                // not kill the dyadic parameters.
                //
                // unfortunately here that is not the case: a huge learning
                // rate will not cause s(h) to increase much, but it will 
                // cause the decay factor to increase greatly

                // dyadic case

                float actualeta = (adaptive) ? eta * rsqrtf (xnorm)
                                             : eta / xnorm;
                float xtx = (adaptive) ? adaptivexnorm : xnorm;
                float norm = (anormplusbnorm + xtx);

                // first try small s approximation

                float smalls;

                if (fabsf (actualeta * norm) < 1e-5)
                  {
                    smalls = (adotb + wdotx - y) * actualeta;
                  }
                else
                  {
                    smalls = (adotb + wdotx - y) * 
                             (1.0f - expf (-actualeta * norm)) / norm;
                  }

                if (fabsf (smalls) < 0.05f)
                  {
                    ++n_smalls;

                    return std::make_pair (smalls, 1.0f);
                  }
                else
                  {
                    //
                    // 
                    // s'[h] == (p (s (h)) - y)
                    

                    DFunc dfunc (adotb, anormplusbnorm, wdotx, xtx, lambda, y);

                    boost::optional<float> s = 
                      dormand_prince_4_5 (dfunc,
                                          0.0f,
                                          0.0f,
                                          actualeta,
                                          1e-2f,
                                          actualeta / 1000.0f,
                                          actualeta,
                                          100);

                    if (s)
                      {
                        return std::make_pair (*s, expf (-actualeta * lambda));
                      }
                    else
                      {
                        ++skipped_update;

                        return std::make_pair (0.0f, 1.0f);
                      }
                  }
              }
          }

        virtual std::ostream& 
        report (std::ostream& o,
                LogLevel      level)
          {
            if (n > 0 && level <= INFO)
              {
                double var = (sumysq / n) - (sumy / n) * (sumy / n);

                o << "best constant = " << (sumy / n) << std::endl;
                o << "best constant loss = " << 0.5 * var << std::endl;
              }

            if (n_smalls > 0 && level <= WARN)
              {
                o << "WARNING: small s approximation utilized " 
                  << n_smalls 
                  << " times (" 
                  << ceil
                       (1000.0 * static_cast<double> (n_smalls) / 
                        static_cast<double> (total_updates)) / 10.0
                  << " %)"
                  << std::endl;
              }

            if (skipped_update > 0 && level <= WARN)
              {
                o << "WARNING: skipped " << skipped_update
                  << " updates ("
                  << ceil
                       (1000.0 * static_cast<double> (skipped_update) / 
                        static_cast<double> (total_updates)) / 10.0
                  << " %)"
                  << " due to loss of precision "
                  << std::endl;
              }

            return o;
          }
    };

  class SquaredClippedLoss : public Loss
    {
      private:
        static LossFactory factory;

        SquaredLoss squared_loss;
        boost::optional<float> minlabel;
        boost::optional<float> maxlabel;

        void
        update_bounds (float y)
          {
            if (! maxlabel || y > *maxlabel)
              {
                maxlabel = y;
              }

            if (! minlabel || y < *minlabel)
              {
                minlabel = y;
              }
          }

        float
        clip (float p)
          {
            return (p > *maxlabel) ? *maxlabel :
                   (p < *minlabel) ? *minlabel : p;
          }

      public:
        explicit
        SquaredClippedLoss () { }

        explicit
        SquaredClippedLoss (boost::optional<float> min, 
                            boost::optional<float> max) 
          : minlabel (min),
            maxlabel (max)
          {
          }

        virtual float loss (float p,
                            float y)
          {
            update_bounds (y);
            return squared_loss.loss (clip (p), y);
          }

        virtual float dldp (float p,
                            float y)
          {
            update_bounds (y);
            return squared_loss.dldp (clip (p), y);
          }

        virtual std::pair<float, float>
        invariant_update (float p,
                          float wdotx,
                          float adotb,
                          float y,
                          float eta,
                          float anormplusbnorm,
                          float xnorm,
                          bool  adaptive,
                          float adaptivexnorm,
                          float lambda)
          {
            update_bounds (y);
            return squared_loss.invariant_update (clip (p), 
                                                  wdotx,
                                                  adotb,
                                                  y, 
                                                  eta, 
                                                  anormplusbnorm,
                                                  xnorm,
                                                  adaptive,
                                                  adaptivexnorm,
                                                  lambda);
          }

        virtual std::ostream& 
        report (std::ostream& o,
                LogLevel      level)
          {
            if (minlabel && level <= INFO)
              {
                o << "minimum label encountered = " << *minlabel << std::endl;
              }

            if (maxlabel && level <= INFO)
              {
                o << "maximum label encountered = " << *maxlabel << std::endl;
              }

            return squared_loss.report (o, level);
          }
    };
}

#endif // __SQUARED_LOSS_HH__
